package org.ricks.gateway.cluster;

import org.ricks.ioc.utils.Logger;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 负载均衡器 - 管理活跃的 Broker 节点
 */
public class LoadBalancer {

    // 存储所有活跃的 Broker (key: clusterAddress)
    private static final Map<String, Broker> activeBrokers = new ConcurrentHashMap<>();
    private static final AtomicInteger counter = new AtomicInteger();

    /**
     * 添加 Broker 节点
     */
    public static void addBroker(Broker broker) {
        if (broker == null || broker.getClusterAddress() == null) {
            Logger.warn("Attempted to add invalid broker: {}", broker);
            return;
        }

        activeBrokers.put(broker.getClusterAddress(), broker);
        printStatus();
    }

    /**
     * 移除 Broker 节点
     */
    public static void removeBroker(Broker broker) {
        if (broker == null || broker.getClusterAddress() == null) {
            return;
        }

        activeBrokers.remove(broker.getClusterAddress());
        printStatus();
    }

    /**
     * 标记节点为可疑状态
     */
    public static void markSuspect(Broker broker) {
        if (broker == null || broker.getClusterAddress() == null) {
            return;
        }

        Broker existing = activeBrokers.get(broker.getClusterAddress());
        if (existing != null) {
            existing.setSuspect(true);
            Logger.warn("Marked broker as suspect: {}", broker.getClusterAddress());
            printStatus();
        }
    }

    /**
     * 获取下一个可用 Broker (轮询算法)
     */
    public static Optional<Broker> nextBroker() {
        if (activeBrokers.isEmpty()) {
            return Optional.empty();
        }

        List<Broker> healthyBrokers = new ArrayList<>();
        List<Broker> suspectBrokers = new ArrayList<>();

        // 分类健康节点和可疑节点
        for (Broker broker : activeBrokers.values()) {
            if (broker.isSuspect()) {
                suspectBrokers.add(broker);
            } else {
                healthyBrokers.add(broker);
            }
        }

        // 优先选择健康节点
        List<Broker> candidates = !healthyBrokers.isEmpty() ? healthyBrokers : suspectBrokers;
        Broker selected = candidates.get(counter.getAndIncrement() % candidates.size());

        Logger.debug("Selected broker: {} (Suspect: {})",
                selected.getClusterAddress(), selected.isSuspect());

        return Optional.of(selected);
    }

    /**
     * 打印当前集群状态
     */
    private static void printStatus() {

        StringBuilder sb = new StringBuilder("\n=== Cluster Status ===");
        sb.append("\nActive Brokers (").append(activeBrokers.size()).append("):");

        activeBrokers.values().forEach(broker -> {
            sb.append("\n - ").append(broker.getId())
                    .append("@").append(broker.getClusterAddress())
                    .append(broker.isSuspect() ? " [SUSPECT]" : "");
        });

        sb.append("\n======================");
        Logger.info(sb.toString());
    }
}
